# -*- coding: utf-8 -*-
import datetime
import os
import sqlite3

from math import floor

from flask import Flask

from flask_sqlalchemy import SQLAlchemy

from get_k_data import get_tdx_stock_minutedata, sh_dir, sz_dir, tdx_dir

import tushare as ts


app = Flask(__name__)
basedir = os.path.abspath(os.path.dirname(__file__))
fullfilename = os.path.join(basedir, '..\\data.sqlite')
print(basedir)
print(fullfilename)
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///' + fullfilename
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = True
db = SQLAlchemy(app)

INIT_CASH = 1000000
fee_rate = 0.00025


class Strategy(db.Model):
    # 定义策略名称、费率、初始资金、起始、结束日期
    __tablename__ = 'strategy'
    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(db.String(40))
    # Excel出现在表头的名称
    title = db.Column(db.String(40))
    fee_rate = db.Column(db.Float)
    init_cash = db.Column(db.Integer)
    start_date = db.Column(db.String(40))
    end_date = db.Column(db.String(40))

    def __init__(self, name, title,
                 fee_rate=fee_rate,
                 init_cash=INIT_CASH,
                 start_date='2017-12-01',
                 end_date='2099-12-31'):
        # 手续费率默认万分之2.5，初始资金100万
        self.name = name
        self.title = title
        self.fee_rate = fee_rate
        self.init_cash = init_cash
        self.start_date = start_date
        self.end_date = end_date

    def __repr__(self):
        return '%s(%r ,%r, %r)' % (self.__class__.__name__, self.id,
                                   self.name, self.title)


# 收盘持仓表，将每天收盘的持仓记录在此
class Holding(db.Model):
    __tablename__ = 'holding'
    id = db.Column(db.Integer, primary_key=True)
    # 日期时间字符串形式
    dt = db.Column(db.String(40))
    name = db.Column(db.String(20))
    stockid = db.Column(db.String(40))
    price = db.Column(db.Integer)
    amount = db.Column(db.Integer)

    sid = db.Column(db.Integer, db.ForeignKey('strategy.id'))
    strategy = db.relationship('Strategy',
                               backref=db.backref('holding',
                                                  lazy='dynamic'))

    def __init__(self, dt, name, stockid, price, amount, strategy):
        self.dt = dt
        self.name = name
        self.stockid = stockid
        self.price = price * 1000
        self.amount = amount
        self.strategy = strategy

    def __repr__(self):
        return '%s(%r ,%r, %r, %r, %r)' % (self.__class__.__name__,
                                           self.dt,
                                           self.name,
                                           self.stockid,
                                           self.price,
                                           self.amount)


# 动态持仓表，买入操作后，放入此表
class DynamicHolding(db.Model):
    __tablename__ = 'dynamic_holding'
    id = db.Column(db.Integer, primary_key=True)
    # 日期时间字符串形式
    dt = db.Column(db.String(40))
    name = db.Column(db.String(20))
    stockid = db.Column(db.String(40))
    price = db.Column(db.Float)
    amount = db.Column(db.Integer)

    sid = db.Column(db.Integer, db.ForeignKey('strategy.id'))
    strategy = db.relationship('Strategy',
                               backref=db.backref('dynamic_holding',
                                                  lazy='dynamic'))

    def __init__(self, dt, name, stockid, price, amount, strategy):
        self.dt = dt
        self.name = name
        self.stockid = stockid
        self.price = price
        self.amount = amount
        self.strategy = strategy

    def __repr__(self):
        return '%s(%r ,%r, %r, %r, %r)' % (self.__class__.__name__,
                                           self.dt,
                                           self.name,
                                           self.stockid,
                                           self.price,
                                           self.amount)


class DailyUpdate(db.Model):
    # 每天策略输入表格
    __tablename__ = 'daily_update'
    id = db.Column(db.Integer, primary_key=True)
    date = db.Column(db.String(40))
    stockid = db.Column(db.String(20))
    # 策略ID
    sid = db.Column(db.Integer, db.ForeignKey('strategy.id'))
    strategy = db.relationship('Strategy',
                               backref=db.backref('daily_update',
                                                  lazy='dynamic'))

    def __init__(self, date, stockid, strategy):
        self.date = date
        self.stockid = stockid
        self.strategy = strategy

    def __repr__(self):
        return '%s(%r ,%r, %r)' % (self.__class__.__name__,
                                   self.sid,
                                   self.date,
                                   self.stockid)


class TradingRecord(db.Model):
    # 模拟交易记录
    __tablename__ = 'trading_record'
    id = db.Column(db.Integer, primary_key=True)
    date = db.Column(db.String(20))
    stockid = db.Column(db.String(20))
    price = db.Column(db.Float)
    # 数量
    amount = db.Column(db.Integer)
    # 成交额
    trans = db.Column(db.Float)
    # 手续费
    fee = db.Column(db.Float)
    sid = db.Column(db.Integer, db.ForeignKey('strategy.id'))
    strategy = db.relationship('Strategy',
                               backref=db.backref('trading_record',
                                                  lazy='dynamic'))

    def __init__(self, date, stockid, price, amount, trans,
                 fee, strategy):
        self.date = date
        self.stockid = stockid
        self.price = price
        self.amount = amount
        self.trans = trans
        self.fee = fee
        self.strategy = strategy

    def __repr__(self):
        return '%s(%r ,%r, %r, %r)' % (self.__class__.__name__,
                                       self.sid, self.date,
                                       self.price, self.amount)


class DailyAsset(db.Model):
    # 模拟交易记录
    __tablename__ = 'daily_asset'
    id = db.Column(db.Integer, primary_key=True)
    date = db.Column(db.String(20))
    cash = db.Column(db.Float)
    holding = db.Column(db.Float)
    asset = db.Column(db.Float)

    sid = db.Column(db.Integer, db.ForeignKey('strategy.id'))
    strategy = db.relationship('Strategy',
                               backref=db.backref('daily_asset',
                                                  lazy='dynamic'))

    def __init__(self, date, cash, holding, asset, strategy):
        self.date = date
        self.cash = cash
        self.holding = holding
        self.asset = asset
        self.strategy = strategy

    def __repr__(self):
        return '%s(%r ,%r, %r, %r, %r)' % (self.__class__.__name__,
                                           self.sid, self.date, self.cash,
                                           self.holding, self.asset)


def add_strategy():
    # 创建新User对象:
    new_strategy = Strategy(name='rg_min', title='RG MIN')
    # 添加到session:
    db.session.add(new_strategy)
    # 提交即保存到数据库:
    db.session.commit()
    # 关闭session:
    db.session.close()


def get_strategy():
    # strategy = session.query(Strategy).all()
    print(db.session.query(Strategy).all())


def get_strategy2():
    conn = sqlite3.connect("../data.sqlite")
    sql = (r"SELECT * FROM strategy")
    results = conn.execute(sql).fetchall()
    print(results)


def get_datelist():
    strategies = Strategy.query.order_by(Strategy.id).all()
    mkt_data_dict = {}
    # log = []
    #预设早上交易时间
    preset_morning_time = "10:00"
    #预设下午交易时间
    preset_afternoon_time = "14:00"
    for strategy in strategies:
        #清除动态持仓表
        dynamic_holding = DynamicHolding.query.all()
        for holding in dynamic_holding:
            db.session.delete(holding)
            db.session.commit()
            print('策略开始，清除动态持仓表 DynamicHolding。')
            log_file('策略开始，清除动态持仓表 DynamicHolding。')
        dates = []
        date_list = db.session.query(DailyUpdate.date).filter_by(
            strategy=strategy).group_by(
            DailyUpdate.date).all()
        fee_rate = strategy.fee_rate
        cash = strategy.init_cash
        print('init_cash=%s' % cash)
        log_file('init_cash=%s' % cash)
        if strategy.id in [1, 3]:
            preset_time = preset_morning_time
        else:
            preset_time = preset_afternoon_time
        print('strategy_id=%s, preset_time=%s' % (strategy.id, preset_time))
        log_file('strategy_id=%s, preset_time=%s' % (strategy.id, preset_time))
        for obj in date_list:
            dt = datetime.datetime.strptime(obj[0], "%Y-%m-%d")
            dates.append(dt)
        #需按日期datetime排序，不能按字符串排序
        dates.sort()
        # print(dates)
        for dt in dates:
            date = '%s-%s-%s' % (dt.year, dt.month, dt.day)
            print('正在处理策略 %s，日期为%s的数据。' % (strategy.id, date))
            log_file('正在处理策略 %s，日期为%s的数据。' % (strategy.id, date))
            #每天根据买入股票，按收盘价累计
            holding_asset = 0
            #卖出原有持仓
            dynamic_holding_list = DynamicHolding.query.all()
            for holding in dynamic_holding_list:
                #匹配行情记录，计算卖出价格
                holding_date, holding_time = holding.dt.split(' ')
                print('holding_date=%s, holding_time=%s' % (holding_date,
                                                            holding_time))
                stockid = holding.stockid
                log_file('len(mkt_data_dict)=%s' % len(mkt_data_dict))
                if stockid in mkt_data_dict:
                    history = mkt_data_dict[stockid]
                    print('缓存中没有 %s 的行情，现已加入。' % stockid)
                    log_file('缓存中没有 %s 的行情，现已加入。' % stockid)
                else:
                    tdx_filename = get_tdx_filename(stockid)
                    history = get_tdx_stock_minutedata(tdx_filename, stockid)
                    mkt_data_dict[stockid] = history
                    print('缓存中存在 %s 的行情，无需再次读取。' % stockid)
                    log_file('缓存中存在 %s 的行情，无需再次读取。' % stockid)
                print('len(history)=%s' % len(history))
                i = 0
                found_date = False
                selling_asset = 0
                while i < len(history) - 1:
                    record = history[i]
                    record2 = history[i + 1]
                    # log_file('record:%s' % record)
                    # log_file('record2:%s' % record2)
                    history_date, history_time = record[1].split(' ')
                    history_date2, history_time2 = record2[1].split(' ')
                    #找到对应操作日期和时间
                    if (holding_date == history_date and
                       holding_date != history_date2):
                        print('找到持仓 %s 对应的卖出日期：%s' % (
                            holding, history_date2))
                        log_file('找到持仓 %s 对应的卖出日期：%s' % (
                            holding, history_date2))
                        found_date = True
                    if found_date:
                        selling_asset = -1
                        # 找到买入时间的前一根K线，9:55和13:55
                        if history_time2 == preset_time:
                            print('找到持仓 %s 对应的卖出时间：%s %s' % (
                                holding, history_date, history_time))
                            log_file('找到持仓 %s 对应的卖出时间：%s %s' % (
                                holding, history_date, history_time))
                            price = round(record[2], 2)
                            amount = holding.amount
                            #删除对应持仓
                            db.session.delete(holding)
                            db.session.commit()
                            #TODO 此处需要补充插入交易记录
                            trans = round(price * amount * 100, 2)
                            fee = round(trans * fee_rate, 2)
                            trading_record = TradingRecord.query.filter(
                                TradingRecord.stockid == stockid).filter(
                                TradingRecord.date == record[1]).filter(
                                TradingRecord.price == price).filter(
                                TradingRecord.amount == amount).first()
                            if trading_record:
                                trading_record.stockid = stockid
                                trading_record.date = record[1]
                                trading_record.price = price
                                trading_record.amount = 0 - amount
                                trading_record.trans = round(
                                    price * amount * 100, 2)
                                trading_record.fee = fee
                                print('交易记录表的记录已存在，现已更新：%s' % trading_record)
                                log_file('交易记录表的记录已存在，现已更新：%s' % trading_record)
                            else:
                                trading_record = TradingRecord(
                                    record[1], stockid, price, 0 - amount,
                                    trans, fee, strategy)
                                db.session.add(trading_record)
                                print('增加交易记录表的记录：%s' % trading_record)
                                log_file('增加交易记录表的记录：%s' % trading_record)
                            db.session.commit()
                            #更新资产变化
                            selling_asset = round(price * amount * 100, 2)
                            old_cash = cash
                            cash = round(cash + selling_asset - fee, 2)
                            print(('卖出后获得 %s 的现金，支出 %s 手续费，现金从 %s '
                                  '变为 %s。' % (selling_asset, fee, old_cash, cash)))
                            log_file(('卖出后获得 %s 的现金，支出 %s 手续费，现金从 %s '
                                     '变为 %s。' % (selling_asset, fee, old_cash, cash)))
                            break
                    i = i + 1
                if i == len(history) - 1:
                    print('没有找到持仓 %s 对应的历史行情记录：%s %s' % (
                        holding, history_date, history_time))
                    log_file('没有找到持仓 %s 对应的历史行情记录：%s %s' % (
                        holding, history_date, history_time))
            #买入新的持仓
            record_count = DailyUpdate.query.filter_by(date=date).count()
            print('策略 %s 在日期 %s 共有 %s 条记录。' % (strategy, date, record_count))
            records = DailyUpdate.query.filter_by(date=date).all()
            for record in records:
                print('%s 当天需要买入的记录：%s' % (date, record))
                log_file('%s 当天需要买入的记录：%s' % (date, record))
            cash_unit = floor(cash / record_count)
            print('总现金为%s, 需要买入%s个股票，每份买入金额为%s' % (
                cash, record_count, cash_unit))
            log_file('总现金为%s, 需要买入%s个股票，每份买入金额为%s' % (
                cash, record_count, cash_unit))
            daily_updates = DailyUpdate.query.filter_by(
                strategy=strategy, date=date).all()
            for daily_update in daily_updates:
                log_file('根据策略内容买入：%s' % daily_update)
                print('根据策略内容买入：%s' % daily_update)
                date_str = daily_update.date
                stockid = daily_update.stockid
                log_file('len(mkt_data_dict)=%s' % len(mkt_data_dict))
                if stockid in mkt_data_dict:
                    history = mkt_data_dict[stockid]
                    print('缓存中没有 %s 的行情，现已加入。' % stockid)
                    log_file('缓存中没有 %s 的行情，现已加入。' % stockid)
                else:
                    tdx_filename = get_tdx_filename(stockid)
                    history = get_tdx_stock_minutedata(tdx_filename, stockid)
                    mkt_data_dict[stockid] = history
                    print('缓存中存在 %s 的行情，无需再次读取。' % stockid)
                    log_file('缓存中存在 %s 的行情，无需再次读取。' % stockid)
                    log_file('len(mkt_data_dict)=%r' % len(mkt_data_dict))
                # print('len(history)= %s' % len(history))
                # log_file('history: %s' % history)
                i = 0
                is_buying = False
                price = -1
                amount = -1
                print('len(history)=%s ' % len(history))
                log_file('len(history)=%s ' % len(history))
                while i < len(history):
                    record = history[i]
                    # print('len(record[%s]): %s' % (i, len(record)))
                    history_date, history_time = record[1].split(' ')
                    # log_file(('history_date=%s, history_time=%s, dt_str=%s'
                    #         '' % (history_date, history_time, dt_str)))
                    #找到对应操作日期和时间
                    if (history_date == date_str and
                       history_time == preset_time):
                        is_buying = True
                        price = round(record[2], 2)
                        amount = round(cash_unit / price / 100)
                        name = ts.get_stock_basics().ix[stockid][0]
                        dt_str = ('%s %s' % (date_str, preset_time))
                        holding = DynamicHolding.query.filter_by(
                            dt=date, stockid=stockid, strategy=strategy).first()
                        if holding:
                            holding.price = price
                            holding.amount = amount
                            print('动态持仓记录已存在？？已更新：%s' % holding)
                            log_file('动态持仓记录已存在？？已更新：%s' % holding)
                        else:
                            holding = DynamicHolding(dt_str, name, stockid,
                                                     price, amount, strategy)
                            db.session.add(holding)
                            print('增加动态持仓记录：%s' % holding)
                            log_file('增加动态持仓记录：%s' % holding)
                        db.session.commit()
                        #增加交易记录
                        trans = round(price * amount * 100, 2)
                        fee = round(trans * fee_rate, 2)
                        trading_record = TradingRecord.query.filter(
                            TradingRecord.stockid == stockid,
                            TradingRecord.date == record[1],
                            TradingRecord.price == price,
                            TradingRecord.amount == amount).first()
                        if trading_record:
                            trading_record.stockid = stockid
                            trading_record.date = record[1]
                            trading_record.price = price
                            trading_record.amount = amount
                            trading_record.trans = round(
                                price * amount * 100, 2)
                            trading_record.fee = round(
                                trading_record.trans * fee_rate, 2)
                            print('交易记录表的记录已存在，现已更新：%s' % trading_record)
                            log_file('交易记录表的记录已存在，现已更新：%s' % trading_record)
                        else:
                            trading_record = TradingRecord(
                                record[1], stockid, price, amount,
                                trans, fee, strategy)
                            db.session.add(trading_record)
                            print('增加交易记录表的记录：%s' % trading_record)
                            log_file('增加交易记录表的记录：%s' % trading_record)
                            db.session.commit()
                        buying_asset = round(price * amount * 100, 2)
                        old_cash = cash
                        cash = round(cash - buying_asset - fee, 2)
                        print(('买入后支出 %s 的现金，支出 %s 手续费，现金从 %s '
                              '变为 %s。' % (buying_asset, fee, old_cash, cash)))
                        log_file(('买入后支出 %s 的现金，支出 %s 手续费，现金从 %s '
                                 '变为 %s。' % (buying_asset, fee, old_cash,
                                  cash)))
                    #增加收盘持仓记录
                    elif (history_date == date_str and is_buying and
                          history_time == "15:00"):
                        price = round(record[2], 2)
                        name = ts.get_stock_basics().ix[stockid][0]
                        fee = round(price * amount * 100 * fee_rate)
                        buying_asset = round(price * amount * 100, 2)
                        holding_asset = holding_asset + buying_asset
                        holding = Holding.query.filter_by(
                            dt=date_str,
                            stockid=stockid,
                            strategy=strategy).first()
                        if holding:
                            holding.price = price
                            holding.amount = amount
                            print('收盘持仓表的记录已存在，现已更新：%s' % holding)
                            log_file('收盘持仓表的记录已存在，现已更新：%s' % holding)
                        else:
                            holding = Holding(date_str, name, stockid,
                                              price, amount, strategy)
                            db.session.add(holding)
                            print('收盘持仓表的记录不存在，新建记录：%s' % holding)
                            log_file('收盘持仓表的记录不存在，新建记录：%s' % holding)
                        db.session.commit()
                        break
                    i += 1
                if i == len(history):
                    print(('没有找到匹配行情 history_date=%s, history_time=%s, '
                           'dt_str=%s' % (history_date, history_time, dt_str)))
                    log_file(('没有找到匹配行情 history_date=%s, history_time=%s, '
                              'dt_str=%s' % (history_date, history_time,
                                             dt_str)))
            #收盘记录持仓，以及当天资产状况
            asset = round(cash + holding_asset, 2)
            daily_asset = DailyAsset.query.filter_by(
                date=date, strategy=strategy).first()
            if daily_asset:
                daily_asset.cash = round(cash, 2)
                daily_asset.holding = round(holding_asset, 2)
                daily_asset.asset = round(cash + holding_asset, 2)
                print('收盘资产记录已存在，已经进行更新：%s' % daily_asset)
                log_file('收盘资产记录已存在，已经进行更新：%s' % daily_asset)
            else:
                daily_asset = DailyAsset(date, cash, holding_asset, asset, strategy)
                db.session.add(daily_asset)
                print('cash=%s, holding_asset=%s, asset=%s' % (cash, holding_asset, asset))
                print('收盘资产记录不存在，新增：%r' % daily_asset)
                log_file('收盘资产记录不存在，新增：%r' % daily_asset)
            db.session.commit()
    # log_file(log)


def get_tdx_filename(stockid):
    if stockid[0] == '6':
        subdir = sh_dir
        mkt = 'sh'
    else:
        subdir = sz_dir
        mkt = 'sz'
    filename = '%s%s.lc5' % (mkt, stockid)
    tdx_filename = os.path.join(tdx_dir, subdir, filename)
    return tdx_filename


def log_file(log_str):
    str_dt = get_now()
    str_dt2 = datetime.datetime.now().strftime("%Y-%m-%d")
    log_filename = '%s.log' % str_dt2
    with open(log_filename, 'a', encoding='utf-8') as my_logfile:
        log_str = '%s %s' % (str_dt, log_str)
        my_logfile.write(log_str + '\n')


def get_now():
    str_dt = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")
    return str_dt


if __name__ == '__main__':
    # add_strategy()
    get_datelist()
    # get_strategy2()
